package com.hunterai.service;

import com.hunterai.entity.GenerateTask;
import com.hunterai.dao.GenerateTaskRepository;
import com.hunterai.comfyui.ComfyUIService;
import com.hunterai.comfyui.TaskStatusResponse;
import com.hunterai.security.JwtUtil;


import java.time.LocalDateTime;
import java.util.UUID;
import java.util.Map;
import java.util.List;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.Pageable;
// import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.JsonNode;



@Service
@RequiredArgsConstructor
@Slf4j
public class GenerateTaskQueueService {
    private final GenerateTaskRepository taskRepository;
    private final ComfyUIService comfyUIService;
    private final JwtUtil jwtUtil;
    // private final ObjectMapper objectMapper;

    // 用户提交任务
    @Transactional
    public String submitTask(String workflowId, Map<String, Object> params, Integer piority, Long userId) {
        GenerateTask task = new GenerateTask();
        task.setTaskId(UUID.randomUUID().toString());
        task.setWorkflowId(workflowId);
        task.setParams(params);
        task.setStatus("QUEUED");
        task.setPiority(piority == null ? 5 : piority);
        task.setUserId(userId);
        task.setCreatedAt(LocalDateTime.now());
        task.setUpdatedAt(LocalDateTime.now());
        taskRepository.save(task);
        return task.getTaskId();
    }

    // 查询用户的所有任务
    public List<GenerateTask> getUserTasks(Long userId) {
        return taskRepository.findByUserIdOrderByCreatedAtDesc(userId);
    }

    // 根据任务ID和用户ID获取任务（确保用户只能查看自己的任务）
    public GenerateTask getTaskByTaskIdAndUserId(String taskId, Long userId) {
        return taskRepository.findByTaskIdAndUserId(taskId, userId);
    }

    // 定时消费队列（可改为@Async线程）
@Scheduled(fixedDelay = 1000)
@Transactional
public void processQueue() {
    log.debug("Start Queue Process...");
    GenerateTask task = taskRepository.findNextQueuedTaskForUpdate();
    if (task == null) {
        log.debug("No task...");
        return;
    }

    try {
        task.setStatus("PROCESSING");
        task.setStartedAt(LocalDateTime.now());
        task.setUpdatedAt(LocalDateTime.now());
        taskRepository.save(task);

        // 获取参数
        String workflowId = task.getWorkflowId();
        Map<String, Object> params = task.getParams();
        String imageName = params != null ? (String) params.get("image") : null;
        // 从params中移除image参数，因为已经单独提取出来了
        if (params != null) {
            params.remove("image");
        }

        // 调用ComfyUI生成图片，获取prompt_id
        JsonNode generateResponse = comfyUIService.generate(workflowId, imageName, params);
        
        if (!generateResponse.has("prompt_id")) {
            throw new Exception("生成任务创建失败：未返回prompt_id");
        }
        
        String promptId = generateResponse.get("prompt_id").asText();
        
        // 轮询等待任务完成并获取图像URL
        String imagePath = null;
        boolean completed = false;
        int maxAttempts = 600; // 最多等待10分钟(600次 * 1秒)
        int attempts = 0;
        
        while (!completed && attempts < maxAttempts) {
            TaskStatusResponse statusResponse = comfyUIService.checkTaskStatus(promptId);
            
            if (statusResponse.isCompleted()) {
                completed = true;
                if (statusResponse.getImages() != null && !statusResponse.getImages().isEmpty()) {
                    imagePath = statusResponse.getImages().get(0).getUrl();
                }
            } else if ("error".equals(statusResponse.getStatus())) {
                throw new Exception("生成任务执行失败: " + statusResponse.getError());
            } else {
                // 等待1秒后再次检查
                Thread.sleep(1000);
                attempts++;
            }
        }
        
        if (!completed) {
            throw new Exception("生成任务超时");
        }
        
        if (imagePath == null) {
            throw new Exception("未找到生成的图像");
        }
        
        task.setImagePath(imagePath);
        task.setStatus("COMPLETED");
        task.setCompletedAt(LocalDateTime.now());
        task.setUpdatedAt(LocalDateTime.now());
        taskRepository.save(task);
    } catch (Exception e) {
        task.setStatus("FAILED");
        task.setErrorMsg(e.getMessage());
        task.setUpdatedAt(LocalDateTime.now());
        taskRepository.save(task);
        log.error("任务处理失败: {}", e.getMessage());
    }
}

    public GenerateTask getTaskByTaskId(String taskId) {
        return taskRepository.findByTaskId(taskId);
    }

    // 分页获取用户任务历史
    public Page<GenerateTask> getUserTasksWithPagination(Long userId, Pageable pageable) {
        return taskRepository.findByUserIdOrderByCreatedAtDesc(userId, pageable);
    }

    public Page<GenerateTask> findByUserIdAndWorkflowId(Long userId, String workflowId, Pageable pageable) {
        return taskRepository.findByUserIdAndWorkflowId(userId, workflowId, pageable);
    }
}